package org.hong.monkey.util

import org.hong.monkey.Logging

import java.io.File
import java.util.PriorityQueue
import scala.util.{Failure, Success, Try}
import org.apache.hadoop.fs.FileSystem

object ShutdownHookManager extends Logging {
  val DEFAULT_SHUTDOWN_PRIORITY = 100

  val MONKEY_CONTEXT_SHUTDOWN_PRIORITY = 50

  val TEMP_DIR_SHUTDOWN_PRIORITY = 25

  private lazy val shutdownHooks = {
    val manager = new MonkeyShutdownHookManager()
    manager.install()
    manager
  }

  private val shutdownDeletePaths = new scala.collection.mutable.HashSet[String]()

  addShutdownHook(TEMP_DIR_SHUTDOWN_PRIORITY) { () =>
    logInfo("Shutdown hook called")
    shutdownDeletePaths.toArray.foreach { dirPath =>
      try {
        logInfo("Deleting directory " + dirPath)
        Utils.deleteRecursively(new File(dirPath))
      } catch {
        case e: Exception => logError(s"Exception while deleting Monkey temp dir: $dirPath", e)
      }
    }
  }

  def registerShutdownDeleteDir(file: File) {
    val absolutePath = file.getAbsolutePath
    shutdownDeletePaths.synchronized {
      shutdownDeletePaths += absolutePath
    }
  }

  def removeShutdownDeleteDir(file: File) {
    val absolutePath = file.getAbsolutePath()
    shutdownDeletePaths.synchronized {
      shutdownDeletePaths.remove(absolutePath)
    }
  }

  def hasRootAsShutdownDeleteDir(file: File): Boolean = {
    val absolutePath = file.getAbsolutePath()
    val retval = shutdownDeletePaths.synchronized {
      shutdownDeletePaths.exists {
        path => !absolutePath.equals(path) && absolutePath.startsWith(path)
      }
    }
    if (retval)
      logInfo(s"path = $file, alread present as root for deletion.")
    retval
  }

  def inShutdown(): Boolean = {
    try {
      val hook = new Thread {
        override def run() = {}
      }
      Runtime.getRuntime.addShutdownHook(hook)
      Runtime.getRuntime.removeShutdownHook(hook)
    } catch {
      case ise: IllegalStateException => return true
    }
    false
  }

  def addShutdownHook(hook: () => Unit): AnyRef = {
    addShutdownHook(DEFAULT_SHUTDOWN_PRIORITY)(hook)
  }

  def addShutdownHook(priority: Int)(hook: () => Unit): AnyRef = {
    shutdownHooks.add(priority, hook)
  }

  def removeShutdownHook(ref: AnyRef): Boolean = {
    shutdownHooks.remove(ref)
  }
}

private[util] class MonkeyShutdownHookManager {
  private val hooks = new PriorityQueue[MonkeyShutdownHook]()
  @volatile private var shuttingDown = false

  def install(): Unit = {
    val hookTask = new Runnable {
      override def run(): Unit = runAll()
    }
    Try(Utils.classForName("org.apache.hadoop.util.ShutdownHookManager")) match {
      case Success(shmClass) =>
        val fsPriority = classOf[FileSystem]
        .getField("SHUTDOWN_HOOK_PRIORITY")
        .get(null)
        .asInstanceOf[Int]
        val shm = shmClass.getMethod("get").invoke(null)
        shm.getClass().getMethod("addShutdownHook", classOf[Runnable], classOf[Int])
        .invoke(shm, hookTask, Integer.valueOf(fsPriority + 30))

      case Failure(_) =>
        Runtime.getRuntime.addShutdownHook(new Thread(hookTask, "Monkey Shutdown Hook"))
    }
  }

  def runAll(): Unit = {
    shuttingDown = true
    var nextHook: MonkeyShutdownHook = null
    while ({nextHook = hooks.synchronized { hooks.poll() }; nextHook != null}) {
      Try(Utils.logUncaughtExceptions(nextHook.run()))
    }
  }

  def add(priority: Int, hook:() => Unit): AnyRef = {
    hooks.synchronized {
      if (shuttingDown)
        throw new IllegalArgumentException("Shutdown hooks cannot be modified during shutdown.")

      val hookRef = new MonkeyShutdownHook(priority, hook)
      hooks.add(hookRef)
      hookRef
    }
  }

  def remove(ref: AnyRef): Boolean = {
    hooks.synchronized { hooks.remove(ref) }
  }
}

private class MonkeyShutdownHook(private val priority: Int, hook: () => Unit)
  extends Comparable[MonkeyShutdownHook] {
  override def compareTo(other: MonkeyShutdownHook): Int = {
    other.priority - priority
  }

  def run(): Unit = hook()

}
